import os
import argparse
from pathlib import Path
from ultralytics import YOLO


def save_results_as_yolov5_format(results, img_path, save_dir='runs/detect/exp', save_conf=True):
    """
    将YOLOv8的检测结果保存为YOLOv5格式的txt文件

    参数:
    results: YOLOv8模型的推理结果
    img_path: 输入图像的路径
    save_dir: 保存txt文件的目录
    save_conf: 是否在txt文件中保存置信度分数
    """
    # 创建保存目录
    os.makedirs(save_dir, exist_ok=True)

    # 获取图片文件名
    img_name = Path(img_path).name
    txt_name = Path(img_name).with_suffix('.txt')
    txt_path = os.path.join(save_dir, txt_name)

    # 获取检测结果
    boxes = results.boxes  # 边界框信息
    cls = boxes.cls.cpu().numpy()  # 类别
    conf = boxes.conf.cpu().numpy()  # 置信度
    bboxes = boxes.xywhn.cpu().numpy()  # 归一化后的边界框(x,y,w,h)

    # 写入txt文件
    with open(txt_path, 'w') as f:
        for i in range(len(bboxes)):
            x, y, w, h = bboxes[i]
            c = int(cls[i])
            line = f"{c} {x:.6f} {y:.6f} {w:.6f} {h:.6f}"
            if save_conf:
                line += f" {conf[i]:.6f}"
            line += '\n'
            f.write(line)

    print(f"已保存检测结果到: {txt_path}")


def save_visualization(results, img_path, save_dir='runs/detect/exp'):
    """
    保存带检测框的可视化图片

    参数:
    results: YOLOv8模型的推理结果
    img_path: 输入图像的路径
    save_dir: 保存可视化图片的目录
    """
    # 创建保存目录
    os.makedirs(save_dir, exist_ok=True)

    # 获取图片文件名
    img_name = Path(img_path).name

    # 保存带检测框的图片
    # 修复：使用 filename 参数而不是 save_dir
    results.save(filename=os.path.join(save_dir, img_name))
    print(f"已保存可视化结果到: {os.path.join(save_dir, img_name)}")


def main(args):
    # 加载模型
    model = YOLO(args.weights)

    # 检查输入路径是文件还是文件夹
    if os.path.isdir(args.source):
        # 处理文件夹中的所有图像
        image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.webp']
        files = [f for f in os.listdir(args.source)
                 if os.path.isfile(os.path.join(args.source, f))
                 and f.lower().endswith(tuple(image_extensions))]

        if not files:
            print(f"错误: 文件夹 {args.source} 中未找到图像文件")
            return

        for file in files:
            file_path = os.path.join(args.source, file)
            print(f"\n处理文件: {file_path}")

            # 进行推理
            results = model(file_path)

            # 保存为YOLOv5格式
            save_results_as_yolov5_format(results[0], file_path,
                                          save_dir=args.save_dir,
                                          save_conf=args.save_conf)

            # 保存可视化结果
            save_visualization(results[0], file_path,
                               save_dir=args.save_dir_vis)

            # 显示结果（如果指定）
            if args.show:
                results[0].show()

    else:
        # 处理单个文件
        print(f"处理文件: {args.source}")
        results = model(args.source)
        save_results_as_yolov5_format(results[0], args.source,
                                      save_dir=args.save_dir,
                                      save_conf=args.save_conf)

        # 保存可视化结果
        save_visualization(results[0], args.source,
                           save_dir=args.save_dir_vis)

        # 显示结果（如果指定）
        if args.show:
            results[0].show()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="使用YOLOv8模型进行目标检测并保存结果")
    parser.add_argument("--weights", type=str, default="./hand_yolov8s.pt", help="模型权重路径")
    parser.add_argument("--source", type=str, default="./datasets/coco8/images/", help="输入图像或文件夹路径")
    parser.add_argument("--save_dir", type=str, default="./datasets/coco8/labels", help="保存YOLOv5格式标注文件的目录")
    parser.add_argument("--save_dir_vis", type=str, default="./datasets/coco8/visualizations",
                        help="保存可视化图片的目录")
    parser.add_argument("--save_conf", action="store_true", help="是否在标注文件中保存置信度分数")
    parser.add_argument("--show", action="store_true", help="是否显示检测结果")

    args = parser.parse_args()
    main(args)